import rospy
import numpy as np
import gdown
import os
os.environ["QT_QPA_PLATFORM_PLUGIN_PATH"] = "/usr/local/lib/python3.8/dist-packages/cv2/qt/plugins/platforms"
import sys
sys.path.append('/workspace/vlmaps')  # 添加 vlmaps 模块的路径
from std_msgs.msg import String
from geometry_msgs.msg import PointStamped
from nav_msgs.msg import Odometry
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import cv2

#from vlmaps.map.vlmap_builder import VLMapBuilder
from vlmaps.utils.llm_utils import parse_object_goal_instruction
from vlmaps.lseg.modules.models.lseg_net import LSegEncNet
#from vlmaps.utils.clip_utils import load_clip_model, encode_text, find_nearest_grid_feat
from omegaconf import DictConfig
from omegaconf import OmegaConf
import torch
from pathlib import Path
import numpy as np
import hydra
from vlmaps.utils.mapping_utils import(grid_id2base_pos_3d,base_pos2grid_id_3d)
from vlmaps.utils.visualize_utils import(get_heatmap_from_mask_3d,visualize_heatmap_3d,visualize_rgb_map_3d)
from std_msgs.msg import Bool




#from vlmaps.task.habitat_object_nav_task import HabitatObjectNavigationTask
#from vlmaps.robot.habitat_lang_robot import HabitatLanguageRobot
from vlmaps.map.vlmap import VLMap
from vlmaps.utils.llm_utils import parse_object_goal_instruction
from vlmaps.utils.matterport3d_categories import mp3dcat,coco_labels
@hydra.main(
    version_base=None,
    config_path="../../config",
    config_name="map_indexing_cfg.yaml",
)

class SemanticNavNode:
    def __init__(self,cfg: DictConfig):
        rospy.init_node('semantic_nav_node')
        self.cfg = cfg
        # --- 参数读取 ---
        self.map_dir = rospy.get_param('~map_dir','/workspace/vlmaps/data/vlmaps_dataset/pool/vlmap')  # 地图目录（包含vlmaps.h5df）
        self.clip_model_name = rospy.get_param('~clip_model', 'ViT-B/32')
        self.text_topic = rospy.get_param('~goal_instruction_topic', '/semantic_goal_instruction')
        self.goal_pose_topic = rospy.get_param('~goal_pose_topic', '/semantic_nav_goal')
        self.odom_topic = rospy.get_param('~odom_topic', '/Odometry')
        # --- 加载地图 ---
        data_dir = Path(cfg.data_paths.vlmaps_data_dir)
        data_dirs = sorted([x for x in data_dir.iterdir() if x.is_dir()])
        scene_id = cfg.scene_id
        self.decay_rate = cfg.decay_rate
        self.map = VLMap(map_config=cfg.map_config,  # 直接使用Hydra加载的配置
            data_dir=data_dirs[scene_id])
        self.map.load_map(data_dirs[scene_id])
        visualize_rgb_map_3d(self.map.grid_pos, self.map.grid_rgb)
        rospy.loginfo("✅ VLMap Loaded: {} grid points".format(len(self.map.grid_feat)))
        # self.lseg_model, self.lseg_transform, self.crop_size, self.base_size, self.norm_mean, self.norm_std = self._init_lseg()
        self.map.customize_obstacle_map(cfg.map_config.potential_obstacle_names, cfg.map_config.obstacle_names, vis=True)
        self.map._init_clip()
        
        #cat = input("What is your interested category in this scene?")
        
        self.map.init_categories(mp3dcat[1:-1]+coco_labels[1:-1])
        # self.mask = self.map.index_map(cat, with_init_cat=True)
        # visualize_masked_map_3d(self.map.grid_pos,  self.mask, self.map.grid_rgb)
        # heatmap = get_heatmap_from_mask_3d(
        #     self.map.grid_pos, self.mask, cell_size=cfg.params.cs, decay_rate=cfg.decay_rate
        # )
        # visualize_heatmap_3d(self.map.grid_pos, heatmap, self.map.grid_rgb)

        # --- ROS通信 ---
        self.goal_pub = rospy.Publisher(self.goal_pose_topic, PointStamped, queue_size=10)
        self.text_sub = rospy.Subscriber(self.text_topic, String, self.instruction_callback)
        self.odom_sub = rospy.Subscriber(self.odom_topic,Odometry,self.odomtery_callback,queue_size=10)
        self.goal_reached_sub = rospy.Subscriber("/goal_reached", Bool, self.goal_reached_callback)
        #the goal queue
        self.goal_queue = []
        self.goal_reached = True
        
        
        #init
        self.curr_pos_on_map = [500,500]
    def instruction_callback(self, msg):
        # 1. 解析语言指令
        text = msg.data

        rospy.loginfo(f"Received instruction: {text}")
        categories = parse_object_goal_instruction(text)
        rospy.loginfo(f"Parsed categories: {categories}")
        # self.mask = self.map.index_map(categories[0],with_init_cat=True)
        # heatmap = get_heatmap_from_mask_3d(
        #     self.map.grid_pos, self.mask, cell_size=self.map.cs, decay_rate=self.decay_rate
        # )
        # visualize_heatmap_3d(self.map.grid_pos, heatmap, self.map.grid_rgb)
        self.goal_pose = PointStamped()
        self.goal_pose.header.stamp = rospy.Time.now()
        self.goal_pose.header.frame_id = "map"  # 使用正确的坐标系
        self.goal_queue = categories  # 保存全部目标
        self.goal_reached = True      # 标记初始到达状态为True
        self.publish_next_goal()      # 发布第一个目标
        # for cat in categories:
        #     # 2. 文本编码集成在get_nearest_pos中
        #     # 3. 匹配地图中最近的点
        #     pos = self.map.get_nearest_pos(self.curr_pos_on_map, cat)#目标点的位置[x,y]
        #     if pos is not None:
        #         x_position,y_position,z_position= grid_id2base_pos_3d(pos[0],pos[1],0,self.map.cs,self.map.gs)
        #         goal_pose.point.x=x_position
        #         goal_pose.point.y=y_position
        #         goal_pose.point.z=z_position
        #         self.goal_pub.publish(goal_pose)
        #         rospy.loginfo(f"Published goal pose for '{cat}': {goal_pose.point}")
        #         if hasattr(self.map, "obstacles_new_cropped"):
        #             from skimage.graph import route_through_array
        #             obs_map = self.map.obstacles_new_cropped.copy()  # bool: 1=free, 0=obstacle
        #             cost_map = np.where(obs_map, 1.0, 1e9)  # 高代价绕开障碍
        #             vis_map = np.stack([obs_map * 255] * 3, axis=-1).astype(np.uint8)

        #             # 原图偏移：减去 rmin/cmin
        #             goal_row, goal_col = pos[0] - self.map.rmin, pos[1] - self.map.cmin
        #             curr_row, curr_col = self.curr_pos_on_map[0] - self.map.rmin, self.curr_pos_on_map[1] - self.map.cmin

        #             # 圈出当前位置和目标位置
        #             if 0 <= goal_row < vis_map.shape[0] and 0 <= goal_col < vis_map.shape[1]:
        #                 cv2.circle(vis_map, (goal_col, goal_row), radius=5, color=(0, 255, 0), thickness=2)  # 绿色圈目标

        #             if 0 <= curr_row < vis_map.shape[0] and 0 <= curr_col < vis_map.shape[1]:
        #                 cv2.circle(vis_map, (curr_col, curr_row), radius=5, color=(0, 0, 255), thickness=2)  # 红色圈当前位置

        #             # 搜索路径：从当前位置到目标位置
        #             try:
        #                 path, _ = route_through_array(cost_map, (curr_row, curr_col), (goal_row, goal_col), fully_connected=True)
        #                 path = np.array(path)

        #                 for y, x in path:
        #                     cv2.circle(vis_map, (x, y), radius=1, color=(255, 0, 255), thickness=-1)  # 用粉色线标记路径
        #             except Exception as e:
        #                 rospy.logwarn(f"⚠️ Path planning failed: {e}")

        #             cv2.imwrite("map.png", vis_map)
        #     else:
        #         rospy.logwarn(f"⚠️ Category '{cat}' not found in map")
        
    def publish_next_goal(self):
        if self.goal_queue and self.goal_reached:
            cat = self.goal_queue.pop(0)
            pos = self.map.get_nearest_pos(self.curr_pos_on_map, cat)
            if pos is not None:
                x_position,y_position,z_position= grid_id2base_pos_3d(pos[0],pos[1],0,self.map.cs,self.map.gs)
                self.goal_pose.point.x=x_position
                self.goal_pose.point.y=y_position
                self.goal_pose.point.z=z_position
                self.goal_pub.publish(self.goal_pose)
                rospy.loginfo(f"Published goal pose for '{cat}': {self.goal_pose.point}")
                if hasattr(self.map, "obstacles_new_cropped"):
                    from skimage.graph import route_through_array
                    obs_map = self.map.obstacles_new_cropped.copy()  # bool: 1=free, 0=obstacle
                    cost_map = np.where(obs_map, 1.0, 1e9)  # 高代价绕开障碍
                    vis_map = np.stack([obs_map * 255] * 3, axis=-1).astype(np.uint8)

                    # 原图偏移：减去 rmin/cmin
                    goal_row, goal_col = pos[0] - self.map.rmin, pos[1] - self.map.cmin
                    curr_row, curr_col = self.curr_pos_on_map[0] - self.map.rmin, self.curr_pos_on_map[1] - self.map.cmin
                    # 圈出当前位置和目标位置
                    if 0 <= goal_row < vis_map.shape[0] and 0 <= goal_col < vis_map.shape[1]:
                        cv2.circle(vis_map, (goal_col, goal_row), radius=5, color=(0, 255, 0), thickness=2)  # 绿色圈目标

                    if 0 <= curr_row < vis_map.shape[0] and 0 <= curr_col < vis_map.shape[1]:
                        cv2.circle(vis_map, (curr_col, curr_row), radius=5, color=(0, 0, 255), thickness=2)  # 红色圈当前位置

                    # 搜索路径：从当前位置到目标位置
                    try:
                        path, _ = route_through_array(cost_map, (curr_row, curr_col), (goal_row, goal_col), fully_connected=True)
                        path = np.array(path)

                        for y, x in path:
                            cv2.circle(vis_map, (x, y), radius=1, color=(255, 0, 255), thickness=-1)  # 用粉色线标记路径
                    except Exception as e:
                        rospy.logwarn(f"⚠️ Path planning failed: {e}")

                    cv2.imwrite("map.png", vis_map)
                    self.goal_reached = False  # 设置为“未到达”
            else:
                rospy.logwarn(f"⚠️ Category '{cat}' not found in map")
                
    def goal_reached_callback(self, msg: Bool):
        if msg.data:  # 收到 True 表示到达
            rospy.loginfo("✅ Reached goal, publishing next one...")
            self.goal_reached = True
            self.publish_next_goal()


                
    def odomtery_callback(self, msg):
        """
        订阅里程计消息并转换到全局栅格地图坐标系
        
        Args:
            msg: 里程计消息,类型通常为nav_msgs/Odometry
        """
        # 从消息中提取机器人在全局坐标系下的位置
        x_base = msg.pose.pose.position.x
        y_base = msg.pose.pose.position.y
        z_base = msg.pose.pose.position.z
        
        # 调用坐标转换函数，将全局坐标转换为栅格坐标
        # 假设self.map.cs和self.map.gs分别是栅格大小和地图尺寸
        row, col, h = base_pos2grid_id_3d(
            self.map.gs, 
            self.map.cs, 
            x_base, 
            y_base, 
            z_base
        )
        
        # 保存到类变量中，格式为[row, col, h]
        self.curr_pos_on_map = [row, col, h]
        
        # 可选：记录日志信息
        rospy.loginfo(f"当前位置 (全局): x={x_base:.3f}, y={y_base:.3f}, z={z_base:.3f}")
        rospy.loginfo(f"当前位置 (栅格): row={row}, col={col}, h={h}")

    
    def _init_lseg(self):
        crop_size = 480  # 480
        base_size = 520  # 520
        if torch.cuda.is_available():
            self.device = "cuda"
        elif torch.backends.mps.is_available():
            self.device = "mps"
        else:
            self.device = "cpu"
        lseg_model = LSegEncNet("", arch_option=0, block_depth=0, activation="lrelu", crop_size=crop_size)
        model_state_dict = lseg_model.state_dict()
        checkpoint_dir = Path(__file__).resolve().parents[1] / "lseg" / "checkpoints"
        checkpoint_path = checkpoint_dir / "demo_e200.ckpt"
        os.makedirs(checkpoint_dir, exist_ok=True)
        if not checkpoint_path.exists():
            print("Downloading LSeg checkpoint...")
            # the checkpoint is from official LSeg github repo
            # https://github.com/isl-org/lang-seg
            checkpoint_url = "https://drive.google.com/u/0/uc?id=1ayk6NXURI_vIPlym16f_RG3ffxBWHxvb"
            gdown.download(checkpoint_url, output=str(checkpoint_path))

        pretrained_state_dict = torch.load(checkpoint_path, map_location=self.device)
        pretrained_state_dict = {k.lstrip("net."): v for k, v in pretrained_state_dict["state_dict"].items()}
        model_state_dict.update(pretrained_state_dict)
        lseg_model.load_state_dict(pretrained_state_dict)

        lseg_model.eval()
        lseg_model = lseg_model.to(self.device)

        norm_mean = [0.5, 0.5, 0.5]
        norm_std = [0.5, 0.5, 0.5]
        lseg_transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        self.clip_feat_dim = lseg_model.out_c
        return lseg_model, lseg_transform, crop_size, base_size, norm_mean, norm_std


if __name__ == '__main__':
    try:
        node = SemanticNavNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
